
* Helper program to print results
capture program drop print_results
program define print_results
	syntax varlist ,  model_names( str  ) stat_names(str) stat_labs(str) [outfile(str )  hidevars_print(real 0 ) ]

	if "`outfile'" == "" & `hidevars_print' == 0  {

		#d ;
			esttab `model_names' ,
			style(tab)  nonumbers star(* 0.10 ** 0.05 *** 0.01 )
			cells(b(star fmt(%9.3f)) se(par))  collabels(none) depvar
			stats(`stat_names', fmt(a3) labels("`stat_labs'" ))
			drop(o.*, relax) keep(`varlist' , relax) append;
		#d cr
	}
	else if "`outfile'" == "" & `hidevars_print' == 1 {
		#d ;
			esttab `model_names' ,
			style(tab) nonumbers star(* 0.10 ** 0.05 *** 0.01 )
			cells(b(star fmt(%9.3f)) se(par))  collabels(none) depvar
			stats(`stat_names', fmt(a3) labels("`stat_labs'" ))
			drop(o.*, relax) keep( , relax) append;
		#d cr
	}
	else if "`outfile'" != "" & `hidevars_print' == 0 {
		#d ;
			esttab `model_names' using "`outfile'",
			style(tab) nonumbers  star(* 0.10 ** 0.05 *** 0.01 )
			cells(b(star fmt(%9.3f)) se(par))  collabels(none) depvar
			stats(`stat_names', fmt(a3)  labels("`stat_labs'" ))
			drop(o.*, relax) keep(`varlist' , relax) replace;
		#d cr
	}
	else if "`outfile'" != "" & `hidevars_print' == 1 {
		#d ;
			esttab `model_names' using "`outfile'",
			style(tab) nonumbers  star(* 0.10 ** 0.05 *** 0.01 )
			cells(b(star fmt(%9.3f)) se(par))  collabels(none) depvar
			stats(`stat_names', fmt(a3)  labels("`stat_labs'" ))
			drop(o.*, relax) keep( , relax) replace;
		#d cr
	}

end

* Program to run Balance, First Stage, RF and 2SLS regressions and save results.
capture program drop group_p_score
* need the following globals: balance scores enrollment
program define group_p_score


	syntax , endogvars(varlist) instruments(varlist) idvar(varname) group_var(varlist) [ pscores(varlist) outcomes(varlist) balance_vars(varlist) other_controls(varlist fv ) ///
			regression_controls(varlist fv ) outfile(str) cluster(varname) non_offered_mean(real 0) 2sls_residuals(real 0) ///
			non_offered_mean(real 0 ) mean_offer(real 0) model_file(str)  joint(real 0) restriction(str) lin_cntrl(real 0)]
	
	preserve
	// balance
	local j = 1
	if "`restriction'" != "" {
		keep if `restriction'
	}
	if "`pscore'" != "" {
		keep if !inlist(`pscore',0,1)
	}

	

		//IV estimates
		* just for esttab
		local endogmod = ""
		foreach var of local instruments{
		foreach group of local group_var {
			gen `var'_`group' = `var'*`group'
			local endogmod " `endogmod' `var'_`group' "
		}
		}
		local total_count: word count `group_var'
		local test_count = 1
		foreach var of local endogvars{
		foreach group of local group_var {
			gen `var'_`group' = `var'*`group'
			local endogvar " `endogvar' `var'_`group' "
			if `test_count' != `total_count' {
				local test_res "`test_res' `var'_`group' = "
			}
			if `test_count' == `total_count' {
				local test_res "`test_res' `var'_`group'"
			}
			local test_count = `test_count' + 1
		}
		}
		
		foreach var of local pscores{
		foreach group of local group_var {
			gen `var'_`group' = `var'*`group'
			local pscore " `pscore' `var'_`group' "
			
		}
		}
		local estimatesstring ""

		local j = 1
		foreach y of local outcomes {

			display in red "outcome: `y', endog: `endogvar', inst: `endogmod', psc: `pscore', cluster: `cluster'"


			if "`weight'" != ""{
				ivreg2 `y' (`endogvar' = `endogmod'  ) `other_controls' `regression_controls' [aw=`weight'], robust partial(`other_controls' `regression_controls')
			}
			else if "`cluster'" != "" & "`pscore'" != "" {
				ivreghdfe `y' (`endogvar' = `endogmod'  ) `other_controls' `regression_controls' , robust absorb(`pscore') cluster(`cluster') partial(`other_controls' `regression_controls')
			}
			else if "`cluster'" != "" {
				ivreghdfe `y' (`endogvar' = `endogmod'  ) `other_controls' `regression_controls' ,  cluster(`cluster') partial(`other_controls' `regression_controls')
			}
			else if "`pscore'" != "" & `lin_cntrl' == 1 {
display in red "model run: pscore with linear control"
				ivreg2 `y' (`endogvar' = `endogmod'  ) `regression_controls' `pscore', robust partial(`regression_controls') 
			}
			else if "`pscore'" != ""{
				ivreghdfe `y' (`endogvar' = `endogmod'  ) `other_controls' `regression_controls', robust partial(`regression_controls') absorb(`pscore')
			}

			else{
				ivreg2 `y' (`endogvar' = `endogmod'  ) `other_controls' `regression_controls', robust partial(`other_controls' `regression_controls')
				
			}
			
			test `test_res'
			
			qui: cap estadd scalar f = r(F)
			qui: cap estadd scalar p = r(p)
			
			qui: cap estadd scalar sargan = e(j)
			qui: cap estadd scalar sarganp = e(jp)

			qui: cap unique `idvar' if e(sample)
			qui: cap estadd scalar obs = `r(sum)'
			
			if `non_offered_mean' == 1 {
			
				qui: cap unique `idvar' if e(sample) & `endogmod' == 0
				qui: cap estadd scalar n_non_off = `r(sum)'
			
				
				qui: cap su `y' if `endogmod' == 0
				qui: cap estadd scalar non_off = r(mean)
				

			}

			qui: estimates store e`j'
			local estimatesstring " `estimatesstring' e`j' "
			qui: if "`model_file'" != "" estimates save "`model_file'_2sls`y'", replace

			local j = `j' + 1

			if `2sls_residuals' == 1 {

				preserve
					keep student_id res_*
					local outfile_resids : subinstr local outfile ".csv" "_2slsresids.dta", all
					sort student_id
					sa "`outfile_resids'", replace
				restore

			}

		}
		
		print_results  `endogvar' , model_names( `estimatesstring' ) hidevars_print( 0 ) outfile("`outfile'_2sls.csv") ///
		stat_names( "obs non_off n_non_off f p `mean_vars'") stat_labs( `"observations"' `"non-offered mean"' `"non-offered num"' "F test" "p-value" `nonoffered_means' )
		
		// display in red: "endogvar: `endogvar'"

		drop `endogmod'


restore

end



